import cv2 as cv
import numpy as np
import dlib
import math

detector = dlib.get_frontal_face_detector()
predictor = dlib.shape_predictor('../dlib/shape_predictor_68_face_landmarks.dat')


def open_camera(camera_id):
    cap = cv.VideoCapture(camera_id)
    while cap.isOpened():
        ok, frame = cap.read()
        if not ok:
            break
        if camera_id == 0 or camera_id == 1:
            frame = cv.flip(frame, 1, dst=None)
        # 获取面部2d关键点
        size = frame.shape
        image_points = get_2d_points(frame)
        try:
            camera_matrix, dist_coeffs, rotation_vec, translation_vec = get_vector(size, image_points)
            # print('rotation:', rotation_vec)
            # 通过给定的内参和外参计算三维点投影到二维图像平面上的坐标
            (nose_end, jacobian) = cv.projectPoints(np.array([(0.0, 0.0, 1000.0)]),
                                                  rotation_vec, translation_vec,
                                                  camera_matrix, dist_coeffs)
            p1 = (int(image_points[0][0]), int(image_points[0][1]))
            p2 = (int(nose_end[0][0][0]), int(nose_end[0][0][1]))
            cv.line(frame, p1, p2, (255, 0, 0), 2)
        except:
            print('not detect face points')
        cv.imshow('frame', frame)
        c = cv.waitKey(40)
        if c & 0xFF == ord('q'):
            break
    cv.destroyAllWindows()


def get_2d_points(frame):
    # 调用dlib库得到人脸关键点
    gray = cv.cvtColor(frame, cv.COLOR_BGR2GRAY)
    rects = detector(gray, 1)
    # 人脸数len(rects)
    # print(len(rects))
    for i in range(len(rects)):
        landmarks = np.matrix([[p.x, p.y] for p in predictor(frame, rects[i]).parts()])
        image_points = np.array([
            (landmarks[30][0, 0], landmarks[30][0, 1]),  # nose tip
            (landmarks[8][0, 0], landmarks[8][0, 1]),  # chin
            (landmarks[36][0, 0], landmarks[36][0, 1]),  # left eye left corner
            (landmarks[45][0, 0], landmarks[45][0, 1]),  # right eye right corner
            (landmarks[48][0, 0], landmarks[48][0, 1]),  # left mouth corner
            (landmarks[54][0, 0], landmarks[54][0, 1]),  # right mouth corner
        ], dtype='double')
        return image_points


# 根据2d点，3d点，相机内参，畸变参数 获取旋转矩阵和平移矩阵
def get_vector(img_size, image_points):
    # 3d模型坐标
    object_model = np.array([
        (0.0, 0.0, 0.0),  # Nose tip
        (0.0, -330.0, -65.0),  # Chin
        (-225.0, 170.0, -135.0),  # Left eye left corner
        (225.0, 170.0, -135.0),  # Right eye right corner
        (-150.0, -150.0, -125.0),  # Left Mouth corner
        (150.0, -150.0, -125.0)  # Right mouth corner
    ])
    # 焦距focal_length(相机坐标系与图像坐标系之间的距离为焦距f，也即图像坐标系原点与焦点重合)
    focal_length = img_size[1]
    center = (img_size[1] / 2, img_size[0] / 2)
    camera_matrix = np.array(
        [[focal_length, 0, center[0]],
         [0, focal_length, center[1]],
         [0, 0, 1]], dtype="double"
    )
    # camera_matrix = np.load('../camera_parameter/mtx.npy')

    # 相机外参假设为0
    dist_coeffs = np.zeros((4, 1))  # Assuming no lens distortion
    # dist_coeffs = np.load('../camera_parameter/dist.npy')
    _, rotation_vector, translation_vector = cv.solvePnP(objectPoints=object_model, imagePoints=image_points,
                                                         cameraMatrix=camera_matrix, distCoeffs=dist_coeffs,
                                                         flags=cv.SOLVEPNP_ITERATIVE)
    return camera_matrix, dist_coeffs, rotation_vector, translation_vector


# 将旋转矩阵转为欧拉角
# def trans_euler_angle(rotation_vector):
#     # calculate rotation angles
#     theta = cv.norm(rotation_vector, cv.NORM_L2)
#
#     # transformed to quaterniond
#     w = math.cos(theta / 2)
#     x = math.sin(theta / 2) * rotation_vector[0][0] / theta
#     y = math.sin(theta / 2) * rotation_vector[1][0] / theta
#     z = math.sin(theta / 2) * rotation_vector[2][0] / theta
#
#     ysqr = y * y
#     # pitch (x-axis rotation)
#     t0 = 2.0 * (w * x + y * z)
#     t1 = 1.0 - 2.0 * (x * x + ysqr)
#     print('t0:{}, t1:{}'.format(t0, t1))
#     pitch = math.atan2(t0, t1)
#
#     # yaw (y-axis rotation)
#     t2 = 2.0 * (w * y - z * x)
#     if t2 > 1.0:
#         t2 = 1.0
#     if t2 < -1.0:
#         t2 = -1.0
#     yaw = math.asin(t2)
#
#     # roll (z-axis rotation)
#     t3 = 2.0 * (w * z + x * y)
#     t4 = 1.0 - 2.0 * (ysqr + z * z)
#     roll = math.atan2(t3, t4)
#
#     print('pitch:{}, yaw:{}, roll:{}'.format(pitch, yaw, roll))
#
#     # 单位转换：将弧度转换为度
#     Y = int((pitch / math.pi) * 180)
#     X = int((yaw / math.pi) * 180)
#     Z = int((roll / math.pi) * 180)
#
#     return 0, Y, X, Z


if __name__ == '__main__':
    video = '../images/te.tif'
    open_camera(0)
